from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
import numpy as np
from typing import List, Optional, Union, Tuple
from math import sqrt, pi, log, sin, cos, exp


class GaussianCircuit:
    """
    A class for creating quantum circuits that produce measurement outcomes
    following a continuous Gaussian distribution.
    
    This class implements the Box-Muller transform approach to generate
    values from a Gaussian distribution. The Box-Muller transform converts
    uniformly distributed random numbers to a normal distribution.
    """
    
    def __init__(self, num_qubits_per_register: int = 4, mean: float = 0.0, std_dev: float = 1.0):
        """
        Initialize the Gaussian distribution circuit generator.
        
        Args:
            num_qubits_per_register: Number of qubits to use for each register (precision)
            mean: Mean of the Gaussian distribution
            std_dev: Standard deviation of the Gaussian distribution
        """
        self.num_qubits_per_register = num_qubits_per_register
        self.mean = mean
        self.std_dev = std_dev
        
        # Total qubits needed for the circuit
        # We need 2 registers for Box-Muller transform: one for uniform [0,1) and one for [0,2π)
        self.total_qubits = 2 * num_qubits_per_register
        
        self.circuit = None
    
    def generate_circuit(self, measure: bool = True) -> QuantumCircuit:
        """
        Generate a quantum circuit that produces Gaussian distributed measurement outcomes.
        
        Args:
            measure: Whether to include measurement operations in the circuit
            
        Returns:
            A quantum circuit that will produce Gaussian distribution outcomes when measured
        """
        # Create quantum and classical registers
        q_uniform = QuantumRegister(self.num_qubits_per_register, 'q_unif')
        q_angle = QuantumRegister(self.num_qubits_per_register, 'q_ang')
        
        # Create classical registers if measurements are requested
        if measure:
            c_uniform = ClassicalRegister(self.num_qubits_per_register, 'c_unif')
            c_angle = ClassicalRegister(self.num_qubits_per_register, 'c_ang')
            qc = QuantumCircuit(q_uniform, q_angle, c_uniform, c_angle)
        else:
            qc = QuantumCircuit(q_uniform, q_angle)
        
        # Step 1: Prepare uniform distributions
        # For both registers, create uniform superpositions using Hadamard gates
        for i in range(self.num_qubits_per_register):
            qc.h(q_uniform[i])
            qc.h(q_angle[i])
        
        # Step 2: Add entanglement to ensure the distributions are truly random
        # Add CNOT gates between adjacent qubits in each register
        for i in range(self.num_qubits_per_register - 1):
            qc.cx(q_uniform[i], q_uniform[i+1])
            qc.cx(q_angle[i], q_angle[i+1])
        
        # Step 3: Apply rotations to the angle register to spread values in [0,2π)
        # Implement a phase gradient to approximate continuous angle distribution
        for i in range(self.num_qubits_per_register):
            # Each qubit position contributes differently to the final angle
            power = self.num_qubits_per_register - i - 1
            angle = 2 * pi / (2**power)
            qc.p(angle, q_angle[i])
        
        # Add measurements if requested
        if measure:
            qc.measure(q_uniform, c_uniform)
            qc.measure(q_angle, c_angle)
        
        self.circuit = qc
        return qc
    
    def process_results(self, counts: dict) -> List[float]:
        """
        Process the measurement results to extract Gaussian distributed values.
        
        Args:
            counts: Dictionary of measurement outcomes and their frequencies
            
        Returns:
            List of Gaussian distributed values derived from the measurements
        """
        gaussian_values = []
        
        for bitstring, frequency in counts.items():
            # Parse the bitstring to extract uniform and angle values
            # The format depends on how the circuit was measured
            if len(bitstring) == 2 * self.num_qubits_per_register:
                # Split the bitstring in half
                uniform_bits = bitstring[:self.num_qubits_per_register]
                angle_bits = bitstring[self.num_qubits_per_register:]
            else:
                # Likely separate registers were measured into separate bitstrings
                # Format might be like "00 01" for c_unif=00, c_ang=01
                parts = bitstring.split()
                if len(parts) >= 2:
                    uniform_bits = parts[0]
                    angle_bits = parts[1]
                else:
                    continue  # Skip if we can't parse properly
            
            # Convert binary strings to decimal values in [0,1) and [0,2π)
            uniform_value = int(uniform_bits, 2) / (2**self.num_qubits_per_register)
            angle_value = int(angle_bits, 2) * 2 * pi / (2**self.num_qubits_per_register)
            
            # Skip 0 for the uniform value (to avoid log(0))
            if uniform_value == 0:
                uniform_value = 1e-10
            
            # Apply Box-Muller transform: Z = sqrt(-2 ln(U)) * cos(2π*V)
            # Where U and V are uniform random variables, Z is standard normal
            r = sqrt(-2 * log(uniform_value))
            z = r * cos(angle_value)
            
            # Scale by standard deviation and shift by mean
            gaussian_value = z * self.std_dev + self.mean
            
            # Add this value to our results, weighting by frequency
            for _ in range(frequency):
                gaussian_values.append(gaussian_value)
        
        return gaussian_values
    
    def plot_distribution(self, counts: dict, bins: int = 30):
        """
        Plot the distribution of values derived from the measurement counts.
        
        Args:
            counts: Dictionary of measurement outcomes and their frequencies
            bins: Number of bins to use in the histogram
            
        Returns:
            Matplotlib figure showing the distribution
        """
        import matplotlib.pyplot as plt
        
        # Process results to get Gaussian values
        values = self.process_results(counts)
        
        # Create plot
        fig, ax = plt.subplots(figsize=(10, 6))
        
        # Plot histogram of values
        ax.hist(values, bins=bins, density=True, alpha=0.7, label='Quantum Circuit Output')
        
        # Plot the ideal Gaussian curve for comparison
        x = np.linspace(min(values), max(values), 1000)
        y = np.exp(-0.5 * ((x - self.mean) / self.std_dev) ** 2) / (self.std_dev * sqrt(2 * pi))
        ax.plot(x, y, 'r-', label='Ideal Gaussian')
        
        # Add labels and title
        ax.set_xlabel('Value')
        ax.set_ylabel('Probability Density')
        ax.set_title('Quantum Circuit Gaussian Distribution (Box-Muller Transform)')
        ax.legend()
        
        # Return the figure for further customization or saving
        return fig


if __name__ == "__main__":
    # Example usage
    gaussian_circuit = GaussianCircuit(num_qubits_per_register=4, mean=0.0, std_dev=1.0)
    circuit = gaussian_circuit.generate_circuit()
    print("Gaussian distribution circuit using Box-Muller transform:")
    print(circuit.draw())
    
    # This would typically be followed by:
    # 1. Execute the circuit on a simulator or real device
    # 2. Process the results to extract Gaussian values
    # 3. Plot the distribution
    
    # Example code for simulation (commented out):
    """
    from qiskit import Aer, execute
    
    # Run the circuit on a simulator
    simulator = Aer.get_backend('qasm_simulator')
    job = execute(circuit, simulator, shots=10000)
    result = job.result()
    counts = result.get_counts(circuit)
    
    # Process and plot the results
    gaussian_values = gaussian_circuit.process_results(counts)
    fig = gaussian_circuit.plot_distribution(counts)
    fig.savefig('gaussian_distribution.png')
    """ 